from __future__ import annotations

import asyncio
import json
import uuid

from langchain_core.tools import StructuredTool  # noqa: TC002

from lfx.base.agents.utils import maybe_unflatten_dict, safe_cache_get, safe_cache_set
from lfx.base.mcp.util import (
    MCPStdioClient,
    MCPStreamableHttpClient,
    create_input_schema_from_json_schema,
    update_tools,
)
from lfx.custom.custom_component.component_with_cache import ComponentWithCache
from lfx.inputs.inputs import InputTypes  # noqa: TC001
from lfx.io import BoolInput, DropdownInput, McpInput, MessageTextInput, Output
from lfx.io.schema import flatten_schema, schema_to_langflow_inputs
from lfx.log.logger import logger
from lfx.schema.dataframe import DataFrame
from lfx.schema.message import Message
from lfx.services.deps import get_settings_service, get_storage_service, session_scope


def resolve_mcp_config(
    server_name: str,  # noqa: ARG001
    server_config_from_value: dict | None,
    server_config_from_db: dict | None,
) -> dict | None:
    """Resolve MCP server config with proper precedence.

    Resolves the configuration for an MCP server with the following precedence:
    1. Database config (takes priority) - ensures edits are reflected
    2. Config from value/tweaks (fallback) - allows REST API to provide config for new servers

    Args:
        server_name: Name of the MCP server
        server_config_from_value: Config provided via value/tweaks (optional)
        server_config_from_db: Config from database (optional)

    Returns:
        Final config to use (DB takes priority, falls back to value)
        Returns None if no config found in either location
    """
    if server_config_from_db:
        return server_config_from_db
    return server_config_from_value


class MCPToolsComponent(ComponentWithCache):
    schema_inputs: list = []
    tools: list[StructuredTool] = []
    _not_load_actions: bool = False
    _tool_cache: dict = {}
    _last_selected_server: str | None = None  # Cache for the last selected server

    def __init__(self, **data) -> None:
        super().__init__(**data)
        # Initialize cache keys to avoid CacheMiss when accessing them
        self._ensure_cache_structure()

        # Initialize clients with access to the component cache
        self.stdio_client: MCPStdioClient = MCPStdioClient(component_cache=self._shared_component_cache)
        self.streamable_http_client: MCPStreamableHttpClient = MCPStreamableHttpClient(
            component_cache=self._shared_component_cache
        )

    def _ensure_cache_structure(self):
        """Ensure the cache has the required structure."""
        # Check if servers key exists and is not CacheMiss
        servers_value = safe_cache_get(self._shared_component_cache, "servers")
        if servers_value is None:
            safe_cache_set(self._shared_component_cache, "servers", {})

        # Check if last_selected_server key exists and is not CacheMiss
        last_server_value = safe_cache_get(self._shared_component_cache, "last_selected_server")
        if last_server_value is None:
            safe_cache_set(self._shared_component_cache, "last_selected_server", "")

    default_keys: list[str] = [
        "code",
        "_type",
        "tool_mode",
        "tool_placeholder",
        "mcp_server",
        "tool",
        "use_cache",
        "verify_ssl",
    ]

    display_name = "MCP Tools"
    description = "Connect to an MCP server to use its tools."
    documentation: str = "https://docs.langflow.org/mcp-tools"
    icon = "Mcp"
    name = "MCPTools"

    inputs = [
        McpInput(
            name="mcp_server",
            display_name="MCP Server",
            info="Select the MCP Server that will be used by this component",
            real_time_refresh=True,
        ),
        BoolInput(
            name="use_cache",
            display_name="Use Cached Server",
            info=(
                "Enable caching of MCP Server and tools to improve performance. "
                "Disable to always fetch fresh tools and server updates."
            ),
            value=False,
            advanced=True,
        ),
        BoolInput(
            name="verify_ssl",
            display_name="Verify SSL Certificate",
            info=(
                "Enable SSL certificate verification for HTTPS connections. "
                "Disable only for development/testing with self-signed certificates."
            ),
            value=True,
            advanced=True,
        ),
        DropdownInput(
            name="tool",
            display_name="Tool",
            options=[],
            value="",
            info="Select the tool to execute",
            show=False,
            required=True,
            real_time_refresh=True,
        ),
        MessageTextInput(
            name="tool_placeholder",
            display_name="Tool Placeholder",
            info="Placeholder for the tool",
            value="",
            show=False,
            tool_mode=False,
        ),
    ]

    outputs = [
        Output(display_name="Response", name="response", method="build_output"),
    ]

    async def _validate_schema_inputs(self, tool_obj) -> list[InputTypes]:
        """Validate and process schema inputs for a tool."""
        try:
            if not tool_obj or not hasattr(tool_obj, "args_schema"):
                msg = "Invalid tool object or missing input schema"
                raise ValueError(msg)

            flat_schema = flatten_schema(tool_obj.args_schema.schema())
            input_schema = create_input_schema_from_json_schema(flat_schema)
            if not input_schema:
                msg = f"Empty input schema for tool '{tool_obj.name}'"
                raise ValueError(msg)

            schema_inputs = schema_to_langflow_inputs(input_schema)
            if not schema_inputs:
                msg = f"No input parameters defined for tool '{tool_obj.name}'"
                await logger.awarning(msg)
                return []

        except Exception as e:
            msg = f"Error validating schema inputs: {e!s}"
            await logger.aexception(msg)
            raise ValueError(msg) from e
        else:
            return schema_inputs

    async def update_tool_list(self, mcp_server_value=None):
        # Accepts mcp_server_value as dict {name, config} or uses self.mcp_server
        mcp_server = mcp_server_value if mcp_server_value is not None else getattr(self, "mcp_server", None)
        server_name = None
        server_config_from_value = None
        if isinstance(mcp_server, dict):
            server_name = mcp_server.get("name")
            server_config_from_value = mcp_server.get("config")
        else:
            server_name = mcp_server
        if not server_name:
            self.tools = []
            return [], {"name": server_name, "config": server_config_from_value}

        # Check if caching is enabled, default to False
        use_cache = getattr(self, "use_cache", False)

        # Use shared cache if available and caching is enabled
        cached = None
        if use_cache:
            servers_cache = safe_cache_get(self._shared_component_cache, "servers", {})
            cached = servers_cache.get(server_name) if isinstance(servers_cache, dict) else None

        if cached is not None:
            try:
                self.tools = cached["tools"]
                self.tool_names = cached["tool_names"]
                self._tool_cache = cached["tool_cache"]
                server_config_from_value = cached["config"]
            except (TypeError, KeyError, AttributeError) as e:
                # Handle corrupted cache data by clearing it and continuing to fetch fresh tools
                msg = f"Unable to use cached data for MCP Server{server_name}: {e}"
                await logger.awarning(msg)
                # Clear the corrupted cache entry
                current_servers_cache = safe_cache_get(self._shared_component_cache, "servers", {})
                if isinstance(current_servers_cache, dict) and server_name in current_servers_cache:
                    current_servers_cache.pop(server_name)
                    safe_cache_set(self._shared_component_cache, "servers", current_servers_cache)
            else:
                return self.tools, {"name": server_name, "config": server_config_from_value}

        try:
            # Try to fetch from database first to ensure we have the latest config
            # This ensures database updates (like editing a server) take effect
            try:
                from langflow.api.v2.mcp import get_server
                from langflow.services.database.models.user.crud import get_user_by_id
            except ImportError as e:
                msg = (
                    "Langflow MCP server functionality is not available. "
                    "This feature requires the full Langflow installation."
                )
                raise ImportError(msg) from e

            server_config_from_db = None
            async with session_scope() as db:
                if not self.user_id:
                    msg = "User ID is required for fetching MCP tools."
                    raise ValueError(msg)
                current_user = await get_user_by_id(db, self.user_id)

                # Try to get server config from DB/API
                server_config_from_db = await get_server(
                    server_name,
                    current_user,
                    db,
                    storage_service=get_storage_service(),
                    settings_service=get_settings_service(),
                )

            # Resolve config with proper precedence: DB takes priority, falls back to value
            server_config = resolve_mcp_config(
                server_name=server_name,
                server_config_from_value=server_config_from_value,
                server_config_from_db=server_config_from_db,
            )

            if not server_config:
                self.tools = []
                return [], {"name": server_name, "config": server_config}

            # Add verify_ssl option to server config if not present
            if "verify_ssl" not in server_config:
                verify_ssl = getattr(self, "verify_ssl", True)
                server_config["verify_ssl"] = verify_ssl

            _, tool_list, tool_cache = await update_tools(
                server_name=server_name,
                server_config=server_config,
                mcp_stdio_client=self.stdio_client,
                mcp_streamable_http_client=self.streamable_http_client,
            )

            self.tool_names = [tool.name for tool in tool_list if hasattr(tool, "name")]
            self._tool_cache = tool_cache
            self.tools = tool_list

            # Cache the result only if caching is enabled
            if use_cache:
                cache_data = {
                    "tools": tool_list,
                    "tool_names": self.tool_names,
                    "tool_cache": tool_cache,
                    "config": server_config,
                }

                # Safely update the servers cache
                current_servers_cache = safe_cache_get(self._shared_component_cache, "servers", {})
                if isinstance(current_servers_cache, dict):
                    current_servers_cache[server_name] = cache_data
                    safe_cache_set(self._shared_component_cache, "servers", current_servers_cache)

        except (TimeoutError, asyncio.TimeoutError) as e:
            msg = f"Timeout updating tool list: {e!s}"
            await logger.aexception(msg)
            raise TimeoutError(msg) from e
        except Exception as e:
            msg = f"Error updating tool list: {e!s}"
            await logger.aexception(msg)
            raise ValueError(msg) from e
        else:
            return tool_list, {"name": server_name, "config": server_config}

    async def update_build_config(self, build_config: dict, field_value: str, field_name: str | None = None) -> dict:
        """Toggle the visibility of connection-specific fields based on the selected mode."""
        try:
            if field_name == "tool":
                try:
                    # Always refresh tools when cache is disabled, or when tools list is empty
                    # This ensures database edits are reflected immediately when cache is disabled
                    use_cache = getattr(self, "use_cache", False)
                    if len(self.tools) == 0 or not use_cache:
                        try:
                            self.tools, build_config["mcp_server"]["value"] = await self.update_tool_list()
                            build_config["tool"]["options"] = [tool.name for tool in self.tools]
                            build_config["tool"]["placeholder"] = "Select a tool"
                        except (TimeoutError, asyncio.TimeoutError) as e:
                            msg = f"Timeout updating tool list: {e!s}"
                            await logger.aexception(msg)
                            if not build_config["tools_metadata"]["show"]:
                                build_config["tool"]["show"] = True
                                build_config["tool"]["options"] = []
                                build_config["tool"]["value"] = ""
                                build_config["tool"]["placeholder"] = "Timeout on MCP server"
                            else:
                                build_config["tool"]["show"] = False
                        except ValueError:
                            if not build_config["tools_metadata"]["show"]:
                                build_config["tool"]["show"] = True
                                build_config["tool"]["options"] = []
                                build_config["tool"]["value"] = ""
                                build_config["tool"]["placeholder"] = "Error on MCP Server"
                            else:
                                build_config["tool"]["show"] = False

                    if field_value == "":
                        return build_config
                    tool_obj = None
                    for tool in self.tools:
                        if tool.name == field_value:
                            tool_obj = tool
                            break
                    if tool_obj is None:
                        msg = f"Tool {field_value} not found in available tools: {self.tools}"
                        await logger.awarning(msg)
                        return build_config
                    await self._update_tool_config(build_config, field_value)
                except Exception as e:
                    build_config["tool"]["options"] = []
                    msg = f"Failed to update tools: {e!s}"
                    raise ValueError(msg) from e
                else:
                    return build_config
            elif field_name == "mcp_server":
                if not field_value:
                    build_config["tool"]["show"] = False
                    build_config["tool"]["options"] = []
                    build_config["tool"]["value"] = ""
                    build_config["tool"]["placeholder"] = ""
                    build_config["tool_placeholder"]["tool_mode"] = False
                    self.remove_non_default_keys(build_config)
                    return build_config

                build_config["tool_placeholder"]["tool_mode"] = True

                current_server_name = field_value.get("name") if isinstance(field_value, dict) else field_value
                _last_selected_server = safe_cache_get(self._shared_component_cache, "last_selected_server", "")
                server_changed = current_server_name != _last_selected_server

                # Determine if "Tool Mode" is active by checking if the tool dropdown is hidden.
                is_in_tool_mode = build_config["tools_metadata"]["show"]

                # Get use_cache setting to determine if we should use cached data
                use_cache = getattr(self, "use_cache", False)

                # Fast path: if server didn't change and we already have options, keep them as-is
                # BUT only if caching is enabled or we're in tool mode
                existing_options = build_config.get("tool", {}).get("options") or []
                if not server_changed and existing_options:
                    # In non-tool mode with cache disabled, skip the fast path to force refresh
                    if not is_in_tool_mode and not use_cache:
                        pass  # Continue to refresh logic below
                    else:
                        if not is_in_tool_mode:
                            build_config["tool"]["show"] = True
                        return build_config

                # To avoid unnecessary updates, only proceed if the server has actually changed
                # OR if caching is disabled (to force refresh in non-tool mode)
                if (_last_selected_server in (current_server_name, "")) and build_config["tool"]["show"] and use_cache:
                    if current_server_name:
                        servers_cache = safe_cache_get(self._shared_component_cache, "servers", {})
                        if isinstance(servers_cache, dict):
                            cached = servers_cache.get(current_server_name)
                            if cached is not None and cached.get("tool_names"):
                                cached_tools = cached["tool_names"]
                                current_tools = build_config["tool"]["options"]
                                if current_tools == cached_tools:
                                    return build_config
                    else:
                        return build_config
                safe_cache_set(self._shared_component_cache, "last_selected_server", current_server_name)

                # When cache is disabled, clear any cached data for this server
                # This ensures we always fetch fresh data from the database
                if not use_cache and current_server_name:
                    servers_cache = safe_cache_get(self._shared_component_cache, "servers", {})
                    if isinstance(servers_cache, dict) and current_server_name in servers_cache:
                        servers_cache.pop(current_server_name)
                        safe_cache_set(self._shared_component_cache, "servers", servers_cache)

                # Check if tools are already cached for this server before clearing
                cached_tools = None
                if current_server_name and use_cache:
                    servers_cache = safe_cache_get(self._shared_component_cache, "servers", {})
                    if isinstance(servers_cache, dict):
                        cached = servers_cache.get(current_server_name)
                        if cached is not None:
                            try:
                                cached_tools = cached["tools"]
                                self.tools = cached_tools
                                self.tool_names = cached["tool_names"]
                                self._tool_cache = cached["tool_cache"]
                            except (TypeError, KeyError, AttributeError) as e:
                                # Handle corrupted cache data by ignoring it
                                msg = f"Unable to use cached data for MCP Server,{current_server_name}: {e}"
                                await logger.awarning(msg)
                                cached_tools = None

                # Clear tools when cache is disabled OR when we don't have cached tools
                # This ensures fresh tools are fetched after database edits
                if not cached_tools or not use_cache:
                    self.tools = []  # Clear previous tools to force refresh

                # Clear previous tool inputs if:
                # 1. Server actually changed
                # 2. Cache is disabled (meaning tool list will be refreshed)
                if server_changed or not use_cache:
                    self.remove_non_default_keys(build_config)

                # Only show the tool dropdown if not in tool_mode
                if not is_in_tool_mode:
                    build_config["tool"]["show"] = True
                    if cached_tools:
                        # Use cached tools to populate options immediately
                        build_config["tool"]["options"] = [tool.name for tool in cached_tools]
                        build_config["tool"]["placeholder"] = "Select a tool"
                    else:
                        # Show loading state only when we need to fetch tools
                        build_config["tool"]["placeholder"] = "Loading tools..."
                        build_config["tool"]["options"] = []
                    # Force a value refresh when:
                    # 1. Server changed
                    # 2. We don't have cached tools
                    # 3. Cache is disabled (to force refresh on config changes)
                    if server_changed or not cached_tools or not use_cache:
                        build_config["tool"]["value"] = uuid.uuid4()
                else:
                    # Keep the tool dropdown hidden if in tool_mode
                    self._not_load_actions = True
                    build_config["tool"]["show"] = False

            elif field_name == "tool_mode":
                build_config["tool"]["placeholder"] = ""
                build_config["tool"]["show"] = not bool(field_value) and bool(build_config["mcp_server"])
                self.remove_non_default_keys(build_config)
                self.tool = build_config["tool"]["value"]
                if field_value:
                    self._not_load_actions = True
                else:
                    build_config["tool"]["value"] = uuid.uuid4()
                    build_config["tool"]["options"] = []
                    build_config["tool"]["show"] = True
                    build_config["tool"]["placeholder"] = "Loading tools..."
            elif field_name == "tools_metadata":
                self._not_load_actions = False

        except Exception as e:
            msg = f"Error in update_build_config: {e!s}"
            await logger.aexception(msg)
            raise ValueError(msg) from e
        else:
            return build_config

    def get_inputs_for_all_tools(self, tools: list) -> dict:
        """Get input schemas for all tools."""
        inputs = {}
        for tool in tools:
            if not tool or not hasattr(tool, "name"):
                continue
            try:
                flat_schema = flatten_schema(tool.args_schema.schema())
                input_schema = create_input_schema_from_json_schema(flat_schema)
                langflow_inputs = schema_to_langflow_inputs(input_schema)
                inputs[tool.name] = langflow_inputs
            except (AttributeError, ValueError, TypeError, KeyError) as e:
                msg = f"Error getting inputs for tool {getattr(tool, 'name', 'unknown')}: {e!s}"
                logger.exception(msg)
                continue
        return inputs

    def remove_non_default_keys(self, build_config: dict) -> None:
        """Remove non-default keys from the build config."""
        for key in list(build_config.keys()):
            if key not in self.default_keys:
                build_config.pop(key)

    async def _update_tool_config(self, build_config: dict, tool_name: str) -> None:
        """Update tool configuration with proper error handling."""
        if not self.tools:
            self.tools, build_config["mcp_server"]["value"] = await self.update_tool_list()

        if not tool_name:
            return

        tool_obj = next((tool for tool in self.tools if tool.name == tool_name), None)
        if not tool_obj:
            msg = f"Tool {tool_name} not found in available tools: {self.tools}"
            self.remove_non_default_keys(build_config)
            build_config["tool"]["value"] = ""
            await logger.awarning(msg)
            return

        try:
            # Store current values before removing inputs (only for the current tool)
            current_values = {}
            for key, value in build_config.items():
                if key not in self.default_keys and isinstance(value, dict) and "value" in value:
                    current_values[key] = value["value"]

            # Remove ALL non-default keys (all previous tool inputs)
            self.remove_non_default_keys(build_config)

            # Get and validate new inputs for the selected tool
            self.schema_inputs = await self._validate_schema_inputs(tool_obj)
            if not self.schema_inputs:
                msg = f"No input parameters to configure for tool '{tool_name}'"
                await logger.ainfo(msg)
                return

            # Add new inputs to build config for the selected tool only
            for schema_input in self.schema_inputs:
                if not schema_input or not hasattr(schema_input, "name"):
                    msg = "Invalid schema input detected, skipping"
                    await logger.awarning(msg)
                    continue

                try:
                    name = schema_input.name
                    input_dict = schema_input.to_dict()
                    input_dict.setdefault("value", None)
                    input_dict.setdefault("required", True)

                    build_config[name] = input_dict

                    # Preserve existing value if the parameter name exists in current_values
                    if name in current_values:
                        build_config[name]["value"] = current_values[name]

                except (AttributeError, KeyError, TypeError) as e:
                    msg = f"Error processing schema input {schema_input}: {e!s}"
                    await logger.aexception(msg)
                    continue
        except ValueError as e:
            msg = f"Schema validation error for tool {tool_name}: {e!s}"
            await logger.aexception(msg)
            self.schema_inputs = []
            return
        except (AttributeError, KeyError, TypeError) as e:
            msg = f"Error updating tool config: {e!s}"
            await logger.aexception(msg)
            raise ValueError(msg) from e

    async def build_output(self) -> DataFrame:
        """Build output with improved error handling and validation."""
        try:
            self.tools, _ = await self.update_tool_list()
            if self.tool != "":
                # Set session context for persistent MCP sessions using Langflow session ID
                session_context = self._get_session_context()
                if session_context:
                    self.stdio_client.set_session_context(session_context)
                    self.streamable_http_client.set_session_context(session_context)
                exec_tool = self._tool_cache[self.tool]
                tool_args = self.get_inputs_for_all_tools(self.tools)[self.tool]
                kwargs = {}
                for arg in tool_args:
                    value = getattr(self, arg.name, None)
                    if value is not None:
                        if isinstance(value, Message):
                            kwargs[arg.name] = value.text
                        else:
                            kwargs[arg.name] = value

                unflattened_kwargs = maybe_unflatten_dict(kwargs)

                output = await exec_tool.coroutine(**unflattened_kwargs)
                tool_content = []
                for item in output.content:
                    item_dict = item.model_dump()
                    item_dict = self.process_output_item(item_dict)
                    tool_content.append(item_dict)

                if isinstance(tool_content, list) and all(isinstance(x, dict) for x in tool_content):
                    return DataFrame(tool_content)
                return DataFrame(data=tool_content)
            return DataFrame(data=[{"error": "You must select a tool"}])
        except Exception as e:
            msg = f"Error in build_output: {e!s}"
            await logger.aexception(msg)
            raise ValueError(msg) from e

    def process_output_item(self, item_dict):
        """Process the output of a tool."""
        if item_dict.get("type") == "text":
            text = item_dict.get("text")
            try:
                return json.loads(text)
            except json.JSONDecodeError:
                return item_dict
        return item_dict

    def _get_session_context(self) -> str | None:
        """Get the Langflow session ID for MCP session caching."""
        # Try to get session ID from the component's execution context
        if hasattr(self, "graph") and hasattr(self.graph, "session_id"):
            session_id = self.graph.session_id
            # Include server name to ensure different servers get different sessions
            server_name = ""
            mcp_server = getattr(self, "mcp_server", None)
            if isinstance(mcp_server, dict):
                server_name = mcp_server.get("name", "")
            elif mcp_server:
                server_name = str(mcp_server)
            return f"{session_id}_{server_name}" if session_id else None
        return None

    async def _get_tools(self):
        """Get cached tools or update if necessary."""
        mcp_server = getattr(self, "mcp_server", None)
        if not self._not_load_actions:
            tools, _ = await self.update_tool_list(mcp_server)
            return tools
        return []
